from gobang import Gobang, WHITE, EMPTY, BLACK, get_slices

import numpy as np
import sklearn
import random
from collections import defaultdict

class GobangAI(object):
    def __init__(self, gobang):
        self.gobang = gobang

        self.c2 = [np.array([1,1,0,0,0])]
        self.t2 = [np.array([0,1,0,0,1,0]), np.array([0,1,0,1,0,0])]
        self.h2 = [np.array([0,0,1,1,0,0,0])]

        self.m3 = [np.array([1,0,0,1,1]), np.array([1,0,1,0,1]), np.array([1,1,1,0,0])]
        self.t3 = [np.array([0,1,0,1,1,0])]
        self.h3 = [np.array([0,1,1,1,0,0])]

        self.c4 = [np.array([1,1,0,1,1]), np.array([1,1,1,0,1]), np.array([1,1,1,1,0])]
        self.h4 = [np.array([0,1,1,1,1,0])]

        self.h5 = [np.array([1,1,1,1,1])]

        self.patterns = [self.h5, self.h4, self.c4, self.h3, self.t3, self.m3, self.h2, self.t2, self.c2]

        self.k1 = 0.7
        self.k2 = 0.3

    def train(self, direction):
        pass

    def guide(self, player=None):
        gobang = self.gobang
        if player == None or player == gobang.cur_player and not gobang.is_over:
            if len(gobang.history) > 0:
                pos = self.get_step(gobang.cur_player)
                gobang.step(gobang.cur_player, pos)
            else:
                pos = (gobang.board_size/2, gobang.board_size/2)
                gobang.step(gobang.cur_player, pos)

    def get_step(self, player):
        # Copy board.
        board = self.gobang.board.copy()
        depth = 1
        score, pos = self.__search(board, depth, player)
        print score
        return pos

    def __search(self, board, depth, player, alpha=-np.inf, beta=np.inf, pos_set=None, history=[]):

        if pos_set == None:
            board_size = board.shape[0]
            pos_set = set()
            for i in xrange(board_size):
                for j in xrange(board_size):
                    pos = (i,j)
                    if not self.__is_empty(board, pos) \
                       and board[pos] == EMPTY:
                        pos_set.add(pos)

        best_pos = None
        if depth <= 0:
            max_score = -np.inf
            score = 0
            for pos, player_ in history:
                if player == player_:
                    score += self.k1 * self.evaluate(board, pos, player_)
                else:
                    score -= self.k2 * self.evaluate(board, pos, player_)
            for pos in pos_set:
                board[pos] = player
                score += self.k1 * self.evaluate(board, pos, player)
                board[pos] = EMPTY
                if score > max_score:
                    max_score, best_pos = score, pos
            return max_score, best_pos
        else:
            candidates = self.sorted_pos_list(board, pos_set, player)
            best_pos = candidates[0]
            for pos in candidates:
                board[pos] = player
                new_pos_set = self.expanded_pos_set(board, pos_set, pos)
                score = - self.__search(board, depth - 1, -player, -beta, -alpha, new_pos_set, history+[(pos, player)])[0]
                board[pos] = EMPTY
                if score > alpha:
                    alpha = score
                    best_pos = pos
                    if score >= beta:
                        break
            return alpha, best_pos

    def sorted_pos_list(self, board, pos_set, player):
        pos_list = []
        for pos in pos_set:
            score = 0
            board[pos] = player
            score += self.k1 * self.evaluate(board, pos, player)
            board[pos] = -player
            score -= self.k2 * self.evaluate(board, pos, -player)
            board[pos] = EMPTY
            pos_list.append((score, pos))
        pos_list.sort(key=lambda x: x[0], reverse=True)
        return [item[1] for item in pos_list]

    def expanded_pos_set(self, board, pos_set, pos, radius=2):
        pos_set = pos_set.copy()
        pos_set.remove(pos)
        r1, r2 = self.__get_range(board.shape[0], pos, radius)
        for i in xrange(r1.start, r1.stop):
            for j in xrange(r2.start, r2.stop):
                if board[pos] == EMPTY:
                    pos_set.add((i, j))
        return pos_set

    def evaluate(self, board, pos, player):
        if player != board[pos]:
            return 0
        radius = self.gobang.n - 1
        slices = get_slices(board, pos, radius)
        patterns = defaultdict(int)
        for s in slices:
            p = self.__get_pattern(s, player)
            patterns[id(p)] += 1
        if patterns[id(self.h5)] > 0:
            return 10 ** 16
        if patterns[id(self.h4)] > 0 or \
           patterns[id(self.c4)] >= 2:
            return 10 ** 10
        if (patterns[id(self.c4)] > 0 and patterns[id(self.h3)] > 0) or \
           (patterns[id(self.c4)] > 0 and patterns[id(self.t3)] > 0):
            return 10 ** 9
        if patterns[id(self.h3)] >= 2 or patterns[id(self.t3)] >= 2:
            return 10 ** 8
        if (patterns[id(self.m3)] > 0 and patterns[id(self.h3)] > 0) or \
           (patterns[id(self.m3)] > 0 and patterns[id(self.t3)] > 0):
            return 10 ** 7
        if patterns[id(self.c4)] > 0:
            return 10 ** 6
        if patterns[id(self.h3)] > 0:
            return 10 ** 5
        if patterns[id(self.t3)] > 0:
            return 10 ** 4
        if patterns[id(self.h2)] >= 2:
            return 10 ** 3
        if patterns[id(self.m3)] > 0:
            return 10 ** 2
        if patterns[id(self.h2)] > 0:
            return 10 ** 1
        if patterns[id(self.t2)] > 0 or patterns[id(self.c2)] > 0:
            return 1
        return 0
        

    def __get_pattern(self, array, player):
        array = array.copy()
        for i in xrange(len(array)):
            if array[i] == player:
                array[i] = 1
            elif array[i] > 0:
                array[i] = 2
        a_len = len(array)
        for pattern in self.patterns:
            for p in pattern:
                p_len = len(p)
                for i in xrange(a_len - p_len + 1):
                    if (array[i:i+p_len] == p).all() \
                       or (array[i:i+p_len] == p[::-1]).all():
                        return pattern

    def __get_range(self, board_size, pos, radius=2):
        row, col = pos
        u = row - radius if row - radius >= 0 else 0
        d = row + radius + 1 if row + radius < board_size else board_size
        l = col - radius if row - radius >= 0 else 0
        r = col + radius + 1 if col + radius < board_size else board_size
        return slice(u,d), slice(l,r)

    def __is_empty(self, board, pos, radius=2):
        r1, r2 = self.__get_range(board.shape[0], pos, radius)
        b = board[r1, r2]
        return np.sum((b == WHITE) + (b == BLACK)) == 0

if __name__ == '__main__':
    pass